"""
数据加载模块
负责从数据库加载和预处理数据
"""

import pandas as pd
import numpy as np
import sqlite3
import streamlit as st
from src.config import DB_PATH, DB_TABLE_NAME


@st.cache_data
def load_data():
    """
    从数据库加载房屋数据并进行预处理
    
    Returns:
        tuple: (X, y, ids) - 特征矩阵、目标变量、ID列表
    """
    # 连接数据库
    conn = sqlite3.connect(DB_PATH)
    query = f"SELECT * FROM {DB_TABLE_NAME}"
    data = pd.read_sql_query(query, conn)
    conn.close()

    # 提取ID列（如果存在）
    if 'Id' in data.columns:
        ids = data['Id']
    else:
        ids = None

    # 分离特征和目标变量
    y = data['SalePrice']
    X = data.drop(['Id', 'SalePrice'], axis=1, errors='ignore')

    # 数据预处理
    X = preprocess_features(X)

    return X, y, ids


def preprocess_features(X):
    """
    预处理特征数据
    
    Args:
        X (pd.DataFrame): 原始特征数据
        
    Returns:
        pd.DataFrame: 预处理后的特征数据
    """
    # 分离数值型和分类型特征
    num_cols = X.select_dtypes(include=[np.number]).columns
    cat_cols = X.select_dtypes(include=['object']).columns

    # 处理缺失值
    X[num_cols] = X[num_cols].fillna(X[num_cols].median())
    X[cat_cols] = X[cat_cols].fillna('None')

    # 对分类特征进行独热编码
    X = pd.get_dummies(X, drop_first=True)

    return X


def get_feature_info():
    """
    获取特征信息
    
    Returns:
        dict: 特征信息字典
    """
    X, y, ids = load_data()
    
    return {
        'feature_count': len(X.columns),
        'sample_count': len(X),
        'feature_names': list(X.columns),
        'target_name': 'SalePrice',
        'has_ids': ids is not None
    }
